from typing import Generic, Iterable, TypeVar

from llmpt.preprocess import Token, Vocab

V = TypeVar('V', bound=Vocab)


class VocabWithSpecials(Vocab, Generic[V]):
    def __init__(self, specials: Iterable[Token], base: V):
        self._specials = tuple(specials)
        self._base = base

    @property
    def base(self) -> V:
        return self._base

    @property
    def specials(self) -> tuple[Token, ...]:
        return self._specials

    @property
    def _offset(self) -> int:
        return len(self._specials)

    def __len__(self) -> int:
        return len(self._base) + self._offset

    def add_token(self, token: str) -> int:
        return self._base.add_token(token) + self._offset

    def get_token(self, index: int, *, raw: bool = False) -> str:
        if raw:
            return self._base.get_token(index)
        if index < self._offset:
            return self._specials[index]
        return self._base.get_token(index - self._offset)

    def index_token(self, token: str, *, raw: bool = False) -> int:
        if raw:
            return self._base.index_token(token)
        if token in self._specials:
            return self._specials.index(token)
        return self._base.index_token(token) + self._offset


class VocabForNLP(VocabWithSpecials[V]):
    PADDING_TOKEN = "<padding>"
    PADDING_INDEX = 0
    UNKNOWN_TOKEN = "<unknown>"
    UNKNOWN_INDEX = 1
    BEGIN_TOKEN = "<begin>"
    BEGIN_INDEX = 2
    END_TOKEN = "<end>"
    END_INDEX = 3
    PREDICT_TOKEN = "<predict>"
    PREDICT_INDEX = 4
    specials = (PADDING_TOKEN, UNKNOWN_TOKEN, BEGIN_TOKEN, END_TOKEN, PREDICT_TOKEN)

    def __init__(self, base: V):
        super().__init__(self.specials, base)

    def index_token(self, token: str, *, raw: bool = False) -> int:
        try:
            return super().index_token(token, raw=raw)
        except KeyError:
            return self.UNKNOWN_INDEX
